import { useState, useEffect, useRef, useCallback } from 'react'

type Callback = () => void
interface CallbackMap {
  [key: string]: Callback
}

type ObjCallback = (objFn: CallbackMap) => void

type SetValueParam<T> = T | ((prevValue: T) => T)

export function useStateCallback<T>(initialValue: T) {
  const [state, setState] = useState<T>(initialValue)
  const cbRef = useRef<Callback | null>(null)
  const currentFnObjRef = useRef<Record<string, Callback>>({}) // 存储最新的函数引用

  const setCallback = useCallback((objFn: CallbackMap) => {
    for (const key in objFn) {
      currentFnObjRef.current[key] = objFn[key]
    }
  }, [])

  const setStateWithCallback = useCallback(
    (value: SetValueParam<T>, callback?: ObjCallback) => {
      console.log('🚀 ~ callback:', callback)

      setState((prev) => {
        const nextState =
          typeof value === 'function' ? (value as (prevValue: T) => T)(prev) : value

        // 设置要执行的回调
        if (callback) {
          cbRef.current = () => {
            callback?.(currentFnObjRef.current) // 执行最新的函数
            cbRef.current = null
          }
        }

        return nextState
      })
    },
    [],
  )

  useEffect(() => {
    if (cbRef.current) {
      cbRef.current()
    }
  }, [state])

  return [state, setStateWithCallback, setCallback] as const
}

// 使用示例
const ExampleComponent = () => {
  const [count, setCount, setCountCallback] = useStateCallback<number>(0)
  const [result, setResult] = useState<string>('')

  const handleUpdate = () => {
    setCount(count + 1, (fnObj) => {
      console.log('fnObj:', fnObj)
      fnObj.calculateResult?.()
    })
  }

  // 这个函数依赖于 count
  const calculateResult = useCallback(() => {
    setResult(`新值: ${count}`) // 使用最新的 count
  }, [count]) // 依赖于 count，当 count 变化时会生成新的函数

  setCountCallback({ calculateResult })

  return (
    <div className="p-4">
      <div className="mb-4">
        <div className="text-xl">Count: {count}</div>
        <div className="text-gray-600">Result: {result}</div>
      </div>

      <div className="space-y-2">
        <button
          className="block w-full bg-blue-500 text-white p-2 rounded"
          onClick={handleUpdate}
        >
          更新并执行回调
        </button>
      </div>
    </div>
  )
}

export default ExampleComponent
